# BigData course project
# Serial version of training algorithm for SOM
# Main module (plotter for whole SOM, color each neuron according to the
# pre-classification of closest train vector)
#
# Arguments:
#
# <som_file_fin>: output of the learn.py program, containing the trained SOM.
#
# <preclass_file>: pre-classification of train vectors, assuming index to
# be train vector position in trainset; and the last field in the line
# to be an string representing the category code (space as separator)
#
# <colormap_file>: mapping of category codes into colors, using names
# understood my pylab. each line has two columnes (space as separator), first 
# column contains the category code and second one contaings color name/code
#
# <plot_file>: output file to be generated (png), where each neuron in the 
# SOM will be plotted with color corresponding to the pre-classif of its
# closest train vector
#

import sys
from pylab import *
from trainset import TrainSet
from som import SOM
from util import *

if len(sys.argv) != 10:
    log("Usage: %s <dim> <ts_size> <ts_file> <size_x> <size_y> " +
        "<som_file_fin> <pre_class> <color_map> <plotmap_file> ", sys.argv[0])
    sys.exit(1)

dim = int(sys.argv[1])
ts_size = int(sys.argv[2])
ts_file = sys.argv[3]
size_x = int(sys.argv[4])
size_y = int(sys.argv[5])
som_file_fin = sys.argv[6]
preclass_file = sys.argv[7]
colormap_file = sys.argv[8]
plotmap_file = sys.argv[9]

ts = TrainSet(dim, ts_size, ts_file)
som = SOM(ts, 1, size_x, size_y, 0, 0, False, som_file_fin)
get_coords = lambda n: (n.x, n.y)
x, y = zip(*map(get_coords, som.lmap))

get_class = lambda l: l.split()[-1]
preclass = map(get_class, read_all_lines(preclass_file))

get_cmap = lambda l: tuple(l.split())
colormap = dict(map(get_cmap, read_all_lines(colormap_file)))

ts.materialize()
def get_color(n):
    i, _ = find_min(ts.cache, n.w)
    try:
        color = colormap[preclass[i]]
    except KeyError:
        color = colormap["DEF"]
    return color
colors = map(get_color, som.lmap)

fig = figure(1)
ax = fig.add_subplot(111, axisbg="black")
ax.scatter(x, y, s=1, marker="s", color=colors)
fig.canvas.draw()
savefig(plotmap_file, bbox_inches="tight")
